// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/aggregate_companion_signatures.h>

#include <pollux/expression/function_signature.h>
#include <pollux/type/type.h>

namespace kumo::pollux::exec {
    namespace {
        // Add type variables used in `type` to `usedVariables`.
        void addUsedVariablesInType(
            const TypeSignature &type,
            const melon::F14FastMap<std::string, SignatureVariable> &allVariables,
            melon::F14FastMap<std::string, SignatureVariable> &usedVariables) {
            auto iter = allVariables.find(type.baseName());
            if (iter != allVariables.end()) {
                usedVariables.emplace(iter->first, iter->second);
            }
            for (const auto &parameter: type.parameters()) {
                addUsedVariablesInType(parameter, allVariables, usedVariables);
            }
        }
    } // namespace

    melon::F14FastMap<std::string, SignatureVariable> usedTypeVariables(
        const std::vector<TypeSignature> &types,
        const melon::F14FastMap<std::string, SignatureVariable> &allVariables) {
        melon::F14FastMap<std::string, SignatureVariable> usedVariables;
        for (const auto &type: types) {
            addUsedVariablesInType(type, allVariables, usedVariables);
        }
        return usedVariables;
    }

    bool isResultTypeResolvableGivenIntermediateType(
        const AggregateFunctionSignaturePtr &signature) {
        auto &allVariables = signature->variables();
        if (allVariables.empty()) {
            return true;
        }

        std::vector<TypeSignature> intermediateType{signature->intermediateType()};
        std::vector<TypeSignature> resultType{signature->returnType()};
        auto variablesInIntermediate =
                usedTypeVariables(intermediateType, allVariables);
        auto variablesInResult = usedTypeVariables(resultType, allVariables);

        for (const auto &[variable, _]: variablesInResult) {
            if (!variablesInIntermediate.count(variable)) {
                return false;
            }
        }
        return true;
    }

    std::string toSuffixString(const TypeSignature &type) {
        auto name = type.baseName();
        // For primitive and decimal types, return their names.
        if (type.parameters().empty() || isDecimalName(name)) {
            return name;
        }
        auto upperName = boost::algorithm::to_upper_copy(name);
        if (upperName == "ARRAY") {
            return "array_" + toSuffixString(type.parameters()[0]);
        }
        if (upperName == "MAP") {
            return "map_" + toSuffixString(type.parameters()[0]) + "_" +
                   toSuffixString(type.parameters()[1]);
        }
        if (upperName == "ROW") {
            name = "row";
        }
        std::string result = name;
        for (const auto &child: type.parameters()) {
            result += '_';
            result += toSuffixString(child);
        }
        result += "_end";
        result += name;
        return result;
    }

    std::vector<AggregateFunctionSignaturePtr>
    CompanionSignatures::partialFunctionSignatures(
        const std::vector<AggregateFunctionSignaturePtr> &signatures) {
        std::vector<AggregateFunctionSignaturePtr> partialSignatures;
        for (const auto &signature: signatures) {
            if (!isResultTypeResolvableGivenIntermediateType(signature)) {
                continue;
            }
            std::vector<TypeSignature> usedTypes = signature->argumentTypes();
            usedTypes.push_back(signature->intermediateType());
            auto variables = usedTypeVariables(usedTypes, signature->variables());

            partialSignatures.push_back(std::make_shared<AggregateFunctionSignature>(
                /*variables*/ variables,
                              /*returnType*/ signature->intermediateType(),
                              /*intermediateType*/ signature->intermediateType(),
                              /*argumentTypes*/ signature->argumentTypes(),
                              /*constantArguments*/ signature->constantArguments(),
                              /*variableArity*/ signature->variableArity()));
        }
        return partialSignatures;
    }

    std::string CompanionSignatures::partialFunctionName(const std::string &name) {
        return name + "_partial";
    }

    AggregateFunctionSignaturePtr CompanionSignatures::mergeFunctionSignature(
        const AggregateFunctionSignaturePtr &signature) {
        if (!isResultTypeResolvableGivenIntermediateType(signature)) {
            return nullptr;
        }

        std::vector<TypeSignature> usedTypes = {signature->intermediateType()};
        auto variables = usedTypeVariables(usedTypes, signature->variables());
        return std::make_shared<AggregateFunctionSignature>(
            /*variables*/ variables,
                          /*returnType*/ signature->intermediateType(),
                          /*intermediateType*/ signature->intermediateType(),
                          /*argumentTypes*/
                          std::vector<TypeSignature>{signature->intermediateType()},
                          /*constantArguments*/ std::vector<bool>{false},
                          /*variableArity*/ false);
    }

    std::vector<AggregateFunctionSignaturePtr>
    CompanionSignatures::mergeFunctionSignatures(
        const std::vector<AggregateFunctionSignaturePtr> &signatures) {
        return processSignaturesOfDistinctIntermediateTypes(
            signatures, [](const AggregateFunctionSignaturePtr &signature) {
                return mergeFunctionSignature(signature);
            });
    }

    std::string CompanionSignatures::mergeFunctionName(const std::string &name) {
        return name + "_merge";
    }

    bool CompanionSignatures::hasSameIntermediateTypesAcrossSignatures(
        const std::vector<AggregateFunctionSignaturePtr> &signatures) {
        std::unordered_set<TypeSignature> seenTypes;
        for (const auto &signature: signatures) {
            auto normalizedType =
                    normalizeType(signature->intermediateType(), signature->variables());
            if (seenTypes.count(normalizedType)) {
                return true;
            }
            seenTypes.insert(normalizedType);
        }
        return false;
    }

    AggregateFunctionSignaturePtr
    CompanionSignatures::mergeExtractFunctionSignature(
        const AggregateFunctionSignaturePtr &signature) {
        if (!isResultTypeResolvableGivenIntermediateType(signature)) {
            return nullptr;
        }

        std::vector<TypeSignature> usedTypes = {
            signature->intermediateType(), signature->returnType()
        };
        auto variables = usedTypeVariables(usedTypes, signature->variables());
        return std::make_shared<AggregateFunctionSignature>(
            /*variables*/ variables,
                          /*returnType*/ signature->returnType(),
                          /*intermediateType*/ signature->intermediateType(),
                          /*argumentTypes*/
                          std::vector<TypeSignature>{signature->intermediateType()},
                          /*constantArguments*/ std::vector<bool>{false},
                          /*variableArity*/ false);
    }

    std::vector<AggregateFunctionSignaturePtr>
    CompanionSignatures::mergeExtractFunctionSignatures(
        const std::vector<AggregateFunctionSignaturePtr> &signatures) {
        return processSignaturesOfDistinctTypes<AggregateFunctionSignaturePtr>(
            signatures, [](const AggregateFunctionSignaturePtr &signature) {
                return mergeExtractFunctionSignature(signature);
            });
    }

    std::string CompanionSignatures::mergeExtractFunctionNameWithSuffix(
        const std::string &name,
        const TypeSignature &resultType) {
        return name + "_merge_extract_" + toSuffixString(resultType);
    }

    std::string CompanionSignatures::mergeExtractFunctionName(
        const std::string &name) {
        return name + "_merge_extract";
    }

    FunctionSignaturePtr CompanionSignatures::extractFunctionSignature(
        const AggregateFunctionSignaturePtr &signature) {
        if (!isResultTypeResolvableGivenIntermediateType(signature)) {
            return nullptr;
        }

        std::vector<TypeSignature> usedTypes = {
            signature->intermediateType(), signature->returnType()
        };
        auto variables = usedTypeVariables(usedTypes, signature->variables());
        return std::make_shared<FunctionSignature>(
            /*variables*/ variables,
                          /*returnType*/ signature->returnType(),
                          /*argumentTypes*/
                          std::vector<TypeSignature>{signature->intermediateType()},
                          /*constantArguments*/ std::vector<bool>{false},
                          /*variableArity*/ false);
    }

    std::string CompanionSignatures::extractFunctionNameWithSuffix(
        const std::string &name,
        const TypeSignature &resultType) {
        return name + "_extract_" + toSuffixString(resultType);
    }

    std::vector<FunctionSignaturePtr>
    CompanionSignatures::extractFunctionSignatures(
        const std::vector<AggregateFunctionSignaturePtr> &signatures) {
        return processSignaturesOfDistinctTypes<FunctionSignaturePtr>(
            signatures, [](const AggregateFunctionSignaturePtr &signature) {
                return extractFunctionSignature(signature);
            });
    }

    std::string CompanionSignatures::extractFunctionName(const std::string &name) {
        return name + "_extract";
    }

    std::unordered_map<TypeSignature, std::vector<AggregateFunctionSignaturePtr> >
    CompanionSignatures::groupSignaturesByReturnType(
        const std::vector<AggregateFunctionSignaturePtr> &signatures) {
        std::unordered_map<TypeSignature, std::vector<AggregateFunctionSignaturePtr> >
                result;
        for (const auto &signature: signatures) {
            result[signature->returnType()].push_back(signature);
        }
        return result;
    }

    TypeSignature CompanionSignatures::normalizeTypeImpl(
        const TypeSignature &type,
        const melon::F14FastMap<std::string, SignatureVariable> &allVariables,
        melon::F14FastMap<std::string, std::string> &renamedVariables) {
        const auto &baseName = type.baseName();

        // Already renamed variables.
        if (renamedVariables.count(baseName)) {
            return TypeSignature{renamedVariables[baseName], {}};
        }
        // Variables to be renamed in consistent manner.
        if (allVariables.count(baseName)) {
            auto normalizedName = fmt::format("T{}", renamedVariables.size());
            renamedVariables[baseName] = normalizedName;
            return TypeSignature{normalizedName, {}};
        }
        // Primitive types.
        if (type.parameters().empty()) {
            return type;
        }
        // Complex types.
        std::vector<TypeSignature> normalizedParameters;
        for (const auto &param: type.parameters()) {
            normalizedParameters.push_back(
                normalizeTypeImpl(param, allVariables, renamedVariables));
        }
        return TypeSignature{baseName, normalizedParameters};
    }

    TypeSignature CompanionSignatures::normalizeType(
        const TypeSignature &type,
        const melon::F14FastMap<std::string, SignatureVariable> &allVariables) {
        melon::F14FastMap<std::string, std::string> renamedVariables;
        return normalizeTypeImpl(type, allVariables, renamedVariables);
    }
} // namespace kumo::pollux::exec
